import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tushare as ts
from datetime import datetime, timedelta

# 初始化tushare
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


class TradingStrategy:
    def __init__(self):
        self.capital = 1000000  # 初始资金
        self.positions = {}  # 持仓
        self.transaction_cost = 0.0015  # 交易成本(佣金+滑点)
        self.stop_loss = 0.1  # 止损比例10%
        self.take_profit = 0.2  # 止盈比例20%
        self.max_position = 10  # 最大持仓股票数
        self.history = []  # 交易记录
        self.portfolio_value = [self.capital]  # 组合价值历史

    def get_data(self, start_date, end_date):
        """获取股票数据"""
        # 获取交易日历
        trade_cal = pro.trade_cal(exchange='', start_date=start_date, end_date=end_date)
        trade_dates = trade_cal[trade_cal['is_open'] == 1]['cal_date'].tolist()

        # 存储所有数据
        all_data = {}

        for date in trade_dates:
            # 获取当日所有股票数据
            df = pro.daily(trade_date=date)
            df = df.sort_values('vol', ascending=False).head(300)  # 取成交量最大的300只股票

            # 计算技术指标
            df['ma5'] = df.groupby('ts_code')['close'].rolling(5).mean().reset_index(level=0, drop=True)
            df['ma20'] = df.groupby('ts_code')['close'].rolling(20).mean().reset_index(level=0, drop=True)
            df['rsi'] = self.calculate_rsi(df, 14)

            all_data[date] = df

        return all_data

    def calculate_rsi(self, df, window=14):
        """计算RSI指标"""
        delta = df.groupby('ts_code')['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)

        avg_gain = gain.groupby(df['ts_code']).transform(lambda x: x.rolling(window).mean())
        avg_loss = loss.groupby(df['ts_code']).transform(lambda x: x.rolling(window).mean())

        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
        return rsi

    def generate_signals(self, data):
        """生成交易信号"""
        signals = []

        for ts_code, group in data.groupby('ts_code'):
            # 简单的双均线策略
            if group['ma5'].iloc[-1] > group['ma20'].iloc[-1] and group['ma5'].iloc[-2] <= group['ma20'].iloc[-2]:
                signals.append((ts_code, 'buy'))
            elif group['ma5'].iloc[-1] < group['ma20'].iloc[-1] and group['ma5'].iloc[-2] >= group['ma20'].iloc[-2]:
                signals.append((ts_code, 'sell'))

        return signals

    def stock_selection(self, data, date):
        """选股与排序"""
        df = data[date].copy()

        # 筛选条件
        df = df[df['vol'] > df['vol'].median()]  # 成交量高于中位数
        df = df[df['close'] > df['ma20']]  # 价格在20日均线之上

        # 排序标准 - 这里使用RSI和成交量的综合评分
        df['score'] = 0.6 * (1 - df['rsi'] / 100) + 0.4 * (df['vol'] / df['vol'].max())

        # 返回排序后的股票
        return df.sort_values('score', ascending=False).head(self.max_position)

    def execute_trades(self, date, signals, selected_stocks, current_prices):
        """执行交易"""
        # 先处理卖出信号
        for ts_code, signal in signals:
            if signal == 'sell' and ts_code in self.positions:
                # 卖出全部持仓
                position = self.positions.pop(ts_code)
                sell_amount = position['shares'] * current_prices[ts_code] * (1 - self.transaction_cost)
                self.capital += sell_amount

                # 记录交易
                self.history.append({
                    'date': date,
                    'ts_code': ts_code,
                    'action': 'sell',
                    'shares': position['shares'],
                    'price': current_prices[ts_code],
                    'value': sell_amount,
                    'pnl': sell_amount - position['cost']
                })

        # 处理买入信号
        available_capital = self.capital / (self.max_position - len(self.positions))

        for ts_code in selected_stocks['ts_code']:
            if ts_code not in self.positions and len(self.positions) < self.max_position:
                # 计算可以买入的股数
                price = current_prices[ts_code]
                max_shares = int(available_capital / (price * (1 + self.transaction_cost)))

                if max_shares > 0:
                    cost = max_shares * price * (1 + self.transaction_cost)
                    self.capital -= cost

                    # 记录持仓
                    self.positions[ts_code] = {
                        'shares': max_shares,
                        'cost': cost,
                        'entry_price': price,
                        'stop_loss': price * (1 - self.stop_loss),
                        'take_profit': price * (1 + self.take_profit)
                    }

                    # 记录交易
                    self.history.append({
                        'date': date,
                        'ts_code': ts_code,
                        'action': 'buy',
                        'shares': max_shares,
                        'price': price,
                        'value': cost,
                        'pnl': 0
                    })

    def check_stop_loss_take_profit(self, date, current_prices):
        """检查止损止盈条件"""
        to_sell = []

        for ts_code, position in self.positions.items():
            current_price = current_prices.get(ts_code, position['entry_price'])

            # 检查止损
            if current_price <= position['stop_loss']:
                to_sell.append((ts_code, 'stop_loss'))
            # 检查止盈
            elif current_price >= position['take_profit']:
                to_sell.append((ts_code, 'take_profit'))

        return to_sell

    def update_portfolio_value(self, date, current_prices):
        """更新组合价值"""
        position_value = sum(
            pos['shares'] * current_prices.get(ts_code, pos['entry_price'])
            for ts_code, pos in self.positions.items()
        )
        total_value = self.capital + position_value
        self.portfolio_value.append(total_value)

        return total_value

    def backtest(self, start_date, end_date):
        """回测主函数"""
        # 获取数据
        all_data = self.get_data(start_date, end_date)
        dates = sorted(all_data.keys())

        for i in range(1, len(dates)):
            prev_date = dates[i - 1]
            current_date = dates[i]

            # 获取当前价格字典
            current_prices = all_data[current_date].set_index('ts_code')['close'].to_dict()

            # 1. 检查止损止盈
            stop_signals = self.check_stop_loss_take_profit(current_date, current_prices)
            if stop_signals:
                self.execute_trades(current_date, stop_signals, pd.DataFrame(), current_prices)

            # 2. 生成交易信号
            signals = self.generate_signals(all_data[prev_date])

            # 3. 选股与排序
            selected_stocks = self.stock_selection(all_data, prev_date)

            # 4. 执行交易
            self.execute_trades(current_date, signals, selected_stocks, current_prices)

            # 5. 更新组合价值
            self.update_portfolio_value(current_date, current_prices)

            # 打印进度
            if i % 10 == 0:
                print(f"Processing {current_date}, Portfolio Value: {self.portfolio_value[-1]:,.2f}")

    def visualize_results(self):
        """可视化回测结果"""
        plt.figure(figsize=(12, 8))

        # 组合价值曲线
        plt.subplot(2, 2, 1)
        plt.plot(self.portfolio_value, label='Portfolio Value')
        plt.title('Portfolio Value Over Time')
        plt.xlabel('Day')
        plt.ylabel('Value')
        plt.grid(True)
        plt.legend()

        # 每日收益率分布
        returns = pd.Series(self.portfolio_value).pct_change().dropna()
        plt.subplot(2, 2, 2)
        returns.hist(bins=50)
        plt.title('Daily Returns Distribution')
        plt.xlabel('Return')
        plt.ylabel('Frequency')

        # 持仓数量变化
        transactions = pd.DataFrame(self.history)
        if not transactions.empty:
            plt.subplot(2, 2, 3)
            transactions['date'] = pd.to_datetime(transactions['date'])
            positions = transactions.groupby('date')['action'].value_counts().unstack().fillna(0)
            positions[['buy', 'sell']].plot(kind='bar', stacked=True, ax=plt.gca())
            plt.title('Buy/Sell Transactions Over Time')
            plt.xlabel('Date')
            plt.ylabel('Count')

        plt.tight_layout()
        plt.show()

        # 打印汇总统计
        print("\n=== Backtest Summary ===")
        print(f"Initial Capital: {self.portfolio_value[0]:,.2f}")
        print(f"Final Portfolio Value: {self.portfolio_value[-1]:,.2f}")
        print(f"Total Return: {(self.portfolio_value[-1] / self.portfolio_value[0] - 1) * 100:.2f}%")

        if not transactions.empty:
            print(f"\nTotal Trades: {len(transactions)}")
            print(f"Buy Trades: {len(transactions[transactions['action'] == 'buy'])}")
            print(f"Sell Trades: {len(transactions[transactions['action'] == 'sell'])}")
            print(f"Average PnL per Trade: {transactions['pnl'].mean():,.2f}")


# 运行策略
if __name__ == "__main__":
    strategy = TradingStrategy()

    # 设置回测日期范围
    end_date = datetime.now().strftime('%Y%m%d')
    start_date = (datetime.now() - timedelta(days=180)).strftime('%Y%m%d')  # 6个月数据

    # 运行回测
    strategy.backtest(start_date, end_date)

    # 可视化结果
    strategy.visualize_results()